import os
import numpy as np
import tifffile
from PIL import Image
import cv2
from tqdm import tqdm
def convert_tif_to_image(input_dir, output_dir, format='png'):
    """
    将TIF文件转换为常见图像格式
    
    Args:
        input_dir: TIF文件所在目录
        output_dir: 输出图像保存目录
        format: 输出图像格式 ('png' 或 'jpg')
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取所有tif文件
    tif_files = sorted([f for f in os.listdir(input_dir) if f.endswith(('.tif', '.tiff'))])
    
    for tif_file in tqdm(tif_files):
        try:
            # 构建完整路径
            tif_path = os.path.join(input_dir, tif_file)
            
            # 读取TIF文件
            img = tifffile.imread(tif_path)
            
            # 处理16位图像
            if img.dtype == np.uint16:
                img = (img / 65535 * 255).astype(np.uint8)
            
            # 确保图像是8位格式
            if img.dtype != np.uint8:
                img = img.astype(np.uint8)
            
            # 如果是单通道图像，转换为RGB
            if len(img.shape) == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
            
            # 构建输出文件名
            output_filename = os.path.splitext(tif_file)[0] + f'.{format}'
            output_path = os.path.join(output_dir, output_filename)
            
            # 保存图像
            Image.fromarray(img).save(output_path, format=format.upper())
            
            print(f"转换成功: {tif_file} -> {output_filename}")
            # exit()
            
        except Exception as e:
            print(f"处理 {tif_file} 时出错: {e}")

# 使用示例
if __name__ == "__main__":
    # 设置输入和输出目录
    input_dir = "/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Desnow/SRRS-2021-val/val/snow_tif"  # TIF文件所在目录
    output_dir = "/ifs/root/ipa01/101/user_101002/Datasets/Multitask_Datasets/Desnow/SRRS-2021-val/val/snow"    # 转换后图像保存目录
    
    # 转换为PNG格式
    convert_tif_to_image(input_dir, output_dir, format='png')